import os
import json
import shutil
import numpy as np
from pathlib import Path
import random

class DataPreparation:
    def __init__(self):
        self.annotations_dir = os.path.abspath("annotations")
        self.dataset_dir = os.path.abspath("dataset")
        self.processed_dir = os.path.abspath("processed_dataset")
        self.splits = ['train', 'val', 'test']
        self.split_ratio = {'train': 0.8, 'val': 0.1, 'test': 0.1}
        
    def create_directories(self):
        """创建必要的目录结构"""
        for split in self.splits:
            for subdir in ['images', 'labels']:
                path = os.path.join(self.processed_dir, split, subdir)
                os.makedirs(path, exist_ok=True)
    
    def process_annotations(self):
        """处理标注文件，生成YOLO格式的标签"""
        print("开始处理标注数据...")
        
        # 检查目录是否存在
        if not os.path.exists(self.annotations_dir):
            raise FileNotFoundError(f"标注目录不存在: {self.annotations_dir}")
        if not os.path.exists(self.dataset_dir):
            raise FileNotFoundError(f"数据集目录不存在: {self.dataset_dir}")
        
        # 统计信息
        total_processed = 0
        skipped_points_count = 0
        total_points = 0
        error_files = []
        
        # 收集所有有效的标注
        valid_annotations = []
        
        # 遍历所有标注文件
        for annotation_file in os.listdir(self.annotations_dir):
            if not annotation_file.endswith('.json'):
                continue
                
            try:
                with open(os.path.join(self.annotations_dir, annotation_file), 'r', encoding='utf-8') as f:
                    annotation = json.load(f)
                
                # 获取图片信息
                image_file = annotation['image']
                image_path = None
                
                # 在数据集目录中查找图片
                for split in self.splits:
                    possible_path = os.path.join(self.dataset_dir, split, image_file)
                    if os.path.exists(possible_path):
                        image_path = possible_path
                        break
                
                if image_path is None:
                    print(f"警告：找不到图片 {image_file}")
                    continue
                    
                # 处理关键点数据
                keypoints = annotation['keypoints']
                image_size = annotation['image_size']
                
                # 计算边界框
                valid_points = [p for p in keypoints if p is not None]
                if not valid_points:
                    print(f"警告：图片 {image_file} 没有有效关键点")
                    continue
                    
                valid_annotations.append((image_path, image_file, keypoints, image_size))
            except Exception as e:
                print(f"处理标注文件 {annotation_file} 时出错: {str(e)}")
                error_files.append(annotation_file)
                continue
        
        if not valid_annotations:
            raise ValueError("没有找到有效的标注数据")
        
        # 随机打乱数据
        random.shuffle(valid_annotations)
        
        # 计算分割点
        n = len(valid_annotations)
        train_end = int(n * self.split_ratio['train'])
        val_end = train_end + int(n * self.split_ratio['val'])
        
        # 分配数据集
        splits_data = {
            'train': valid_annotations[:train_end],
            'val': valid_annotations[train_end:val_end],
            'test': valid_annotations[val_end:]
        }
        
        # 处理每个分割
        for split, annotations in splits_data.items():
            print(f"\n处理{split}集...")
            for image_path, image_file, keypoints, image_size in annotations:
                try:
                    # 复制图片到处理后的目录
                    target_image_path = os.path.join(self.processed_dir, split, 'images', image_file)
                    shutil.copy2(image_path, target_image_path)
                    
                    # 计算边界框
                    valid_points = [p for p in keypoints if p is not None]
                    x_coords = [p[0] for p in valid_points]
                    y_coords = [p[1] for p in valid_points]
                    
                    x_min, x_max = min(x_coords), max(x_coords)
                    y_min, y_max = min(y_coords), max(y_coords)
                    
                    # 添加一些padding
                    padding = 20
                    x_min = max(0, x_min - padding)
                    y_min = max(0, y_min - padding)
                    x_max = min(image_size[0], x_max + padding)
                    y_max = min(image_size[1], y_max + padding)
                    
                    # 计算边界框的中心点和宽高
                    bbox_width = x_max - x_min
                    bbox_height = y_max - y_min
                    bbox_x_center = (x_min + x_max) / 2
                    bbox_y_center = (y_min + y_max) / 2
                    
                    # 归一化边界框坐标
                    bbox_x_center_norm = bbox_x_center / image_size[0]
                    bbox_y_center_norm = bbox_y_center / image_size[1]
                    bbox_width_norm = bbox_width / image_size[0]
                    bbox_height_norm = bbox_height / image_size[1]
                    
                    # 处理关键点
                    points_data = []
                    for point in keypoints:
                        if point is None:
                            skipped_points_count += 1
                            points_data.extend([0.0, 0.0, 0])  # x, y, visibility=0
                        else:
                            x, y = point
                            # 归一化坐标
                            x_norm = x / image_size[0]
                            y_norm = y / image_size[1]
                            points_data.extend([x_norm, y_norm, 2])  # visibility=2表示可见
                        total_points += 1
                    
                    # 生成YOLO格式的标签文件
                    label_file = os.path.join(
                        self.processed_dir, 
                        split, 
                        'labels', 
                        os.path.splitext(image_file)[0] + '.txt'
                    )
                    
                    # 写入标签文件：类别ID + 边界框 + 归一化的关键点坐标
                    with open(label_file, 'w') as f:
                        # 确保每个关键点都有正确的格式
                        formatted_points = []
                        for i in range(0, len(points_data), 3):
                            x, y, v = points_data[i:i+3]
                            formatted_points.extend([f"{float(x):.6f}", f"{float(y):.6f}", f"{int(v)}"])
                        
                        # 写入格式化的数据，确保有44列
                        line = f"0 {float(bbox_x_center_norm):.6f} {float(bbox_y_center_norm):.6f} {float(bbox_width_norm):.6f} {float(bbox_height_norm):.6f}"
                        for i in range(0, len(formatted_points), 3):
                            line += f" {formatted_points[i]} {formatted_points[i+1]} {formatted_points[i+2]}"
                        
                        parts = line.split()
                        if len(parts) != 44:  # 1(class) + 4(bbox) + 13*3(keypoints) = 44
                            print(f"警告：标签文件 {label_file} 列数不正确：{len(parts)}")
                            continue
                        f.write(line + '\n')
                    
                    total_processed += 1
                except Exception as e:
                    print(f"处理图片 {image_file} 时出错: {str(e)}")
                    error_files.append(image_file)
                    continue
            
            print(f"{split}集处理完成，包含{len(annotations)}张图片")
        
        # 打印统计信息
        print(f"\n数据处理完成！")
        print(f"处理的图片总数：{total_processed}")
        print(f"总关键点数：{total_points}")
        print(f"跳过的关键点数：{skipped_points_count}")
        print(f"跳过率：{(skipped_points_count/total_points*100):.2f}%")
        if error_files:
            print(f"处理失败的文件：{len(error_files)}个")
            for file in error_files:
                print(f"  - {file}")
        
        # 生成数据集配置文件
        self.generate_dataset_config()
    
    def generate_dataset_config(self):
        """生成YOLO格式的数据集配置文件"""
        config = {
            'path': self.processed_dir,  # 使用绝对路径
            'train': 'train/images',
            'val': 'val/images',
            'test': 'test/images',
            'nc': 1,  # 类别数量
            'names': ['pickleball_player'],
            'kpt_shape': [13, 3],  # [关键点数量, 维度(x,y,visible)]
            'flip_idx': [0, 2, 1, 4, 3, 6, 5, 8, 7, 10, 9, 12, 11],  # 水平翻转时关键点对应关系
            'keypoints': [
                '头部', '左肩', '右肩', 
                '左肘', '右肘', '左手腕', '右手腕',
                '左臀', '右臀', '左膝', '右膝',
                '左脚踝', '右脚踝'
            ],
            'connections': [
                [0, 1], [0, 2],  # 头部到肩膀
                [1, 3], [3, 5],  # 左臂
                [2, 4], [4, 6],  # 右臂
                [1, 7], [2, 8],  # 肩膀到臀部
                [7, 9], [9, 11],  # 左腿
                [8, 10], [10, 12]  # 右腿
            ]
        }
        
        with open('dataset.yaml', 'w', encoding='utf-8') as f:
            import yaml
            yaml.dump(config, f, allow_unicode=True, sort_keys=False)
        
        print("\n已生成数据集配置文件：dataset.yaml")

def main():
    try:
        data_prep = DataPreparation()
        data_prep.create_directories()
        data_prep.process_annotations()
        print("\n数据准备完成！接下来可以开始训练模型了。")
    except Exception as e:
        print(f"数据准备过程中出错: {str(e)}")
        raise

if __name__ == "__main__":
    main() 